from trinity.sceneRenderJobInterior import SceneRenderJobInterior
import interiorVisualizations as iVis
import trinity
import blue
import geo2
import math

def CreateWodSceneRenderJobInterior(name = None, stageKey = None):
    newRJ = WodSceneRenderJobInterior()
    if (name is not None):
        newRJ.ManualInit(name)
    else:
        newRJ.ManualInit()
    newRJ.SetMultiViewStage(stageKey)
    return newRJ



class WodSceneRenderJobInterior(SceneRenderJobInterior):
    renderStepOrder = ['UPDATE_SCENE',
     'UPDATE_UI',
     'SET_VIEWPORT',
     'SET_PROJECTION',
     'SET_VIEW',
     'VISIBILITY_QUERY',
     'SET_VISUALIZATION',
     'SET_DEPTH',
     'SET_PREPASS_RT',
     'CLEAR_PREPASS',
     'BEGIN_MANAGED_RENDERING',
     'RENDER_PREPASS',
     'SET_LIGHT_RT',
     'CLEAR_LIGHTS',
     'SET_VAR_LIGHTS_DEPTH',
     'SET_VAR_LIGHTS_PREPASS',
     'RENDER_LIGHTS',
     'RENDER_SSAO',
     'SET_GATHER_RT',
     'SET_VAR_LIGHTS',
     'CLEAR_GATHER_COLOR',
     'SET_SRGB_WRITE_ON',
     'RENDER_GATHER',
     'SET_SRGB_WRITE_OFF',
     'SET_VAR_FRAMEBUFFER_DIMENSIONS',
     'SET_VAR_GATHERFRAMEBUFFER_RTT',
     'SET_DOWNSAMPLE8_RT',
     'DOWN_SAMPLE_8X',
     'SET_DOWNSAMPLE16_RT',
     'DOWN_SAMPLE_2X',
     'SET_VAR_HORIZONTAL8_RT',
     'HORIZONTAL_FILTER_8',
     'SET_VAR_VERTICAL8_RT',
     'VERTICAL_FILTER_8',
     'SET_VAR_HORIZONTAL16_RT',
     'HORIZONTAL_FILTER_16',
     'SET_VAR_VERTICAL16_RT',
     'VERTICAL_FILTER_16',
     'SET_FINAL_RT',
     'POST_PROCESS',
     'END_MANAGED_RENDERING',
     'UPDATE_TOOLS',
     'RENDER_INFO',
     'RENDER_VISUAL',
     'RENDER_TOOLS',
     'RESTORE_DEPTH',
     'RENDER_UI']

    def _ManualInit(self, name = 'WodSceneRenderJobInterior'):
        SceneRenderJobInterior._ManualInit(self)
        self._enablePostProcess = True
        self.gatherFormat = None
        self.gatherTarget = None
        self.scratchA = None
        self.scratchB = None
        self.scratch16A = None
        self.scratch16B = None



    def SetupPostProcessSteps(self):
        sm = trinity.GetShaderManager()
        for s in sm.shaderLibrary:
            if (s.name == 'PostProcess'):
                self.postShader = trinity.Tr2ShaderMaterial()
                self.postShader.highLevelShaderName = 'PostProcess'
                self.postShader.PopulateDefaultParameters()
                lut = self.postShader.parameters['RGBLookup']
                lut.resourcePath = 'res:/graphics/texture/BaseRGBLUT.dds'
                self.AddStep('SET_GATHER_RT', trinity.TriStepSetRenderTarget())

        if (self.postShader is not None):
            self.AddStep('POST_PROCESS', trinity.TriStepRenderFullScreenShader(self.postShader))



    def _CreateBasicRenderSteps(self):
        SceneRenderJobInterior._CreateBasicRenderSteps(self)
        self.AddStep('SET_FINAL_RT', trinity.TriStepSetRenderTarget())
        self.SetupPostProcessSteps()



    def ChooseBufferFormats(self):
        self.depthFormat = trinity.TRIFMT_INTZ
        self.prePassFormat = trinity.TRIFMT_A16B16G16R16
        self.lightAccumulationFormat = trinity.TRIFMT_A8R8G8B8
        self.gatherFormat = trinity.TRIFMT_A8R8G8B8
        customDepthFormat = True
        if not customDepthFormat:
            self.RemoveStep('SET_DEPTH')
            self.RemoveStep('RESTORE_DEPTH')



    def DoReleaseResources(self, level):
        self.renderTargetList = None
        self.backBufferDepthStencil = None
        self.backBufferRenderTarget = None
        self.prePassDepthStencil = None
        self.prePassRenderTarget = None
        self.lightAccumulationTarget = None
        self.gatherTarget = None
        self.scratchA = None
        self.scratchB = None
        self.scratch16A = None
        self.scratch16B = None
        self.SetRenderTargets(None, None, None, None, None, None)



    def SetupOffsetsAndWeights(self, shader):
        o1 = geo2.Vector(-6.2857, -4.2857, -2.3667, -0.4746)
        o2 = geo2.Vector(1.4513, 3.3519, 5.32, 7.3077)
        w1 = geo2.Vector(0.0077, 0.0461, 0.1647, 0.324)
        w2 = geo2.Vector(0.3042, 0.1186, 0.0275, 0.0071)
        p = shader.parameters['FilterWeights1']
        p.value = w1
        p = shader.parameters['FilterWeights2']
        p.value = w2
        p = shader.parameters['FilterOffsets1']
        p.value = o1
        p = shader.parameters['FilterOffsets2']
        p.value = o2



    def SetupBloomShaderParameters(self, width, height):
        scratchWidth = math.floor((width / 8))
        scratchHeight = math.floor((height / 8))
        ds = geo2.Vector(scratchWidth, scratchHeight, (1.0 / scratchWidth), (1.0 / scratchHeight))
        self.AddStep('SET_VAR_SOURCE8_DIMENSIONS', trinity.TriStepSetVariableStore('SourceDimensions', ds))
        scratch16Width = math.floor((scratchWidth / 2))
        scratch16Height = math.floor((scratchHeight / 2))
        d16 = geo2.Vector(scratch16Width, scratch16Height, (1.0 / scratch16Width), (1.0 / scratch16Height))
        self.AddStep('SET_VAR_SOURCE16_DIMENSIONS', trinity.TriStepSetVariableStore('SourceDimensions', d16))
        self.SetupOffsetsAndWeights(self.filterShaderHorizontal8)
        p = self.filterShaderHorizontal8.parameters['FilterStuff']
        p.x = 1.0
        p.y = 0.0
        self.SetupOffsetsAndWeights(self.filterShaderHorizontal16)
        p = self.filterShaderHorizontal8.parameters['FilterStuff']
        p.x = 1.0
        p.y = 0.0
        self.SetupOffsetsAndWeights(self.filterShaderVertical8)
        p = self.filterShaderHorizontal8.parameters['FilterStuff']
        p.x = 0.0
        p.y = 1.0
        self.SetupOffsetsAndWeights(self.filterShaderVertical16)
        p = self.filterShaderHorizontal8.parameters['FilterStuff']
        p.x = 0.0
        p.y = 1.0



    def SetupBloomPingPongBuffers(self, scratchWidth, scratchHeight):
        isw = int(scratchWidth)
        ish = int(scratchHeight)
        if ((self.gatherFormat is not None) and (self.scratchA is None)):
            self.scratchA = device.CreateTexture(isw, ish, 1, trinity.TRIUSAGE_RENDERTARGET, self.gatherFormat, trinity.TRIPOOL_DEFAULT)
        if ((self.gatherFormat is not None) and (self.scratchB is None)):
            self.scratchB = device.CreateTexture(isw, ish, 1, trinity.TRIUSAGE_RENDERTARGET, self.gatherFormat, trinity.TRIPOOL_DEFAULT)
        isw = int((isw / 2))
        ish = int((ish / 2))
        if ((self.gatherFormat is not None) and (self.scratch16A is None)):
            self.scratch16A = device.CreateTexture(isw, ish, 1, trinity.TRIUSAGE_RENDERTARGET, self.gatherFormat, trinity.TRIPOOL_DEFAULT)
        if ((self.gatherFormat is not None) and (self.scratch16B is None)):
            self.scratch16B = device.CreateTexture(isw, ish, 1, trinity.TRIUSAGE_RENDERTARGET, self.gatherFormat, trinity.TRIPOOL_DEFAULT)



    def DoPrepareResources(self):
        if ((not self.enabled) or (not self.canCreateRenderTargets)):
            return 
        if (self.renderTargetList is None):
            device = trinity.device
            self.ChooseBufferFormats()
            if (self.backBufferDepthStencil is None):
                self.backBufferDepthStencil = device.GetDepthStencilSurface()
            if (self.backBufferRenderTarget is None):
                self.backBufferRenderTarget = device.GetRenderTarget()
            width = device.GetPresentParameters()['BackBufferWidth']
            height = device.GetPresentParameters()['BackBufferHeight']
            d = geo2.Vector(width, height, (1.0 / width), (1.0 / height))
            self.AddStep('SET_VAR_FRAMEBUFFER_DIMENSIONS', trinity.TriStepSetVariableStore('FramebufferDimensions', d))
            if ((self.depthFormat is not None) and (self.prePassDepthStencil is None)):
                if (self.depthFormat == self.backBufferDepthStencil.format):
                    self.prePassDepthStencil = self.backBufferDepthStencil
                else:
                    self.prePassDepthStencil = device.CreateTexture(width, height, 1, trinity.TRIUSAGE_DEPTHSTENCIL, self.depthFormat, trinity.TRIPOOL_DEFAULT)
            if ((self.prePassFormat is not None) and (self.prePassRenderTarget is None)):
                self.prePassRenderTarget = device.CreateTexture(width, height, 1, trinity.TRIUSAGE_RENDERTARGET, self.prePassFormat, trinity.TRIPOOL_DEFAULT)
            if ((self.lightAccumulationFormat is not None) and (self.lightAccumulationTarget is None)):
                self.lightAccumulationTarget = device.CreateTexture(width, height, 1, trinity.TRIUSAGE_RENDERTARGET, self.lightAccumulationFormat, trinity.TRIPOOL_DEFAULT)
            if ((self.gatherFormat is not None) and (self.gatherTarget is None)):
                self.gatherTarget = device.CreateTexture(width, height, 1, trinity.TRIUSAGE_RENDERTARGET, self.gatherFormat, trinity.TRIPOOL_DEFAULT)
            noiseTex = trinity.GetNoiseTexture()
            trinity.GetVariableStore().RegisterVariable('NoiseTexture', noiseTex)
            self.renderTargetList = (blue.BluePythonWeakRef(self.backBufferDepthStencil),
             blue.BluePythonWeakRef(self.backBufferRenderTarget),
             blue.BluePythonWeakRef(self.prePassDepthStencil),
             blue.BluePythonWeakRef(self.prePassRenderTarget),
             blue.BluePythonWeakRef(self.lightAccumulationTarget),
             blue.BluePythonWeakRef(self.gatherTarget))
        thingToSet = (x.object for x in self.renderTargetList)
        self.SetRenderTargets(*thingToSet)



    def SetSSAO(self, ssaoStep):
        return 



    def SetRenderTargets(self, backBufferDepthStencil, backBufferRenderTarget, prePassDepthStencil, prePassRenderTarget, lightAccumulationTarget, gatherTarget):
        SceneRenderJobInterior.SetRenderTargets(self, backBufferDepthStencil, backBufferRenderTarget, prePassDepthStencil, prePassRenderTarget, lightAccumulationTarget)
        if (gatherTarget is not None):
            self.SetStepAttr('SET_GATHER_RT', 'target', gatherTarget.GetSurfaceLevel(0))
            self.AddStep('SET_VAR_GATHERFRAMEBUFFER_RTT', trinity.TriStepSetVariableStore('Framebuffer', gatherTarget))
            if (self.postShader is not None):
                trinity.GetVariableStore().RegisterVariable('Framebuffer', gatherTarget)
                self.postShader.BindLowLevelShader(['none'])
        else:
            self.SetStepAttr('SET_GATHER_RT', 'target', None)
        self.UpdatePostProcess(gatherTarget, backBufferRenderTarget)
        self.AddStep('SET_SRGB_WRITE_ON', trinity.TriStepSetRenderState(trinity.D3DRS_SRGBWRITEENABLE, 1))
        self.AddStep('SET_SRGB_WRITE_OFF', trinity.TriStepSetRenderState(trinity.D3DRS_SRGBWRITEENABLE, 0))
        self.AddStep('CLEAR_GATHER_COLOR', trinity.TriStepClear((0, 0, 0.4, 0), None, None))



    def UpdatePostProcess(self, gatherTarget, backbufferTarget):
        if not gatherTarget:
            return 
        if self._enablePostProcess:
            self.SetStepAttr('SET_GATHER_RT', 'target', gatherTarget.GetSurfaceLevel(0))
            self.SetStepAttr('SET_FINAL_RT', 'enabled', True)
            self.SetStepAttr('POST_PROCESS', 'enabled', True)
        else:
            self.SetStepAttr('SET_GATHER_RT', 'target', backbufferTarget)
            self.SetStepAttr('SET_FINAL_RT', 'enabled', False)
            self.SetStepAttr('POST_PROCESS', 'enabled', False)



    def GetEnablePostProcess(self):
        return self._enablePostProcess



    def SetEnablePostProcess(self, isEnabled):
        if (self.gatherTarget is None):
            isEnabled = False
        self._enablePostProcess = isEnabled
        self.UpdatePostProcess(self.gatherTarget, self.backBufferRenderTarget)




